import os
from PIL import Image
import numpy as np
from torch.utils.data import Dataset


class DriveDataset(Dataset):
    def __init__(self, root: str, train: bool,rdd: bool, transforms=None):
        super(DriveDataset, self).__init__()
        self.flag = "training" if train else "test"
        if rdd:
            self.flag = "redundance"
        data_root = os.path.join(root, "SLIVER07", self.flag)                                                                       #拼接训练/测试数据集路径
        assert os.path.exists(data_root), f"path '{data_root}' does not exists."                               #判断文件夹是否存在
        self.transforms = transforms
        img_names = [i for i in os.listdir(os.path.join(data_root, "images")) if i.endswith(".png")]        #遍历存入每个图片的名字
        self.img_list = [os.path.join(data_root, "images", i) for i in img_names]                                  #拼接原图文件路径
        self.manual = [os.path.join(data_root, "1st_manual", i) for i in img_names]                            #拼接标注好的文件的路径

        # check files
        for i in self.manual:
            if os.path.exists(i) is False:
                raise FileNotFoundError(f"file {i} does not exists.")

        self.roi_mask = [os.path.join(data_root, "mask", i) for i in img_names]                                  #拼接mask文件路径

        # check files
        for i in self.roi_mask:
           if os.path.exists(i) is False:
               raise FileNotFoundError(f"file {i} does not exists.")

    def __getitem__(self, idx):
        img = Image.open(self.img_list[idx]).convert('L')                                   #如果图片不是rgb，要改
        #img = np.array(img)                                                                                #改
        manual = Image.open(self.manual[idx]).convert('L')
        manual = np.array(manual) / 255                                                              #前景像素值255/255=1，背景0/255=0
        roi_mask = Image.open(self.roi_mask[idx]).convert('L')
        roi_mask = 255 - np.array(roi_mask)                                                       #让感兴趣的区域像素值为0，不感兴趣的地方像素值为255，方便在计算损失的时候将像素值为255的忽略掉
        mask = np.clip(manual + roi_mask, a_min=0, a_max=255)                        #手工标注和蒙板合在一起后（注意这里的mask和上面的roi_mask不是同一个东西），感兴趣区域前景像素值为1，背景为0，不感兴趣区域为255

        # 这里转回PIL的原因是，transforms中是对PIL数据进行处理
        mask = Image.fromarray(mask)
        #img = Image.fromarray(img)                                                                    #改

        if self.transforms is not None:
            img, mask = self.transforms(img, mask)

        return img, mask

    def __len__(self):
        return len(self.img_list)

    @staticmethod
    def collate_fn(batch):
        images, targets = list(zip(*batch))                                                           #zip() 函数用于将多个可迭代对象的元素进行配对打包，返回一个迭代器，迭代器中的每个元素是一个元组，每个元组代表一个样本的数据，例如 (image1, target1)、(image2, target2) 等。在这里，batch 是一个批处理数据，其中每个元素代表一个样本，可能包含图像和对应的目标（标签）。
                                                                                                                          # 这种操作常用于在训练过程中，将批处理数据中的图像和目标分开，以便于后续的数据处理和模型训练。在分割后，images 列表中存储了所有的图像数据，而 targets 列表中存储了所有的目标数据
        batched_imgs = cat_list(images, fill_value=0)
        batched_targets = cat_list(targets, fill_value=255)
        return batched_imgs, batched_targets

def cat_list(images, fill_value=0):
    max_size = tuple(max(s) for s in zip(*[img.shape for img in images]))        #(3,480,480)
    batch_shape = (len(images),) + max_size                                                     #(4,3,480,480)
    batched_imgs = images[0].new(*batch_shape).fill_(fill_value)                  #根据 images[0] 的数据类型和设备创建一个新的张量，并将其形状设置为 batch_shape。
    for img, pad_img in zip(images, batched_imgs):                                          #将数据打包成一个tensor（batched_imgs），作为输入到网络中的一个batch
        pad_img[..., :img.shape[-2], :img.shape[-1]].copy_(img)                           #这里使用切片 [..., :img.shape[-2], :img.shape[-1]] （img.shape[-2] 和 img.shape[-1] 分别表示 img 张量的高度和宽度）选择了 pad_img 张量中与 img 相同尺寸的部分，并将其内容替换为 img 的内容。
    return batched_imgs

